/**
 * 内存管理系统初始化前期的内存分配工具。
 */

#include <mm/early_alloc.h>
#include <mm/pages.h>
#include <mm/page_flags.h>
#include <cds/bitmap.h>

/* 从页号到页内首个字节的地址需要经历的偏移量。 */
static unsigned long __page_shift = 12;
/* 做内存页相关计算所需的 mask */
static unsigned long __page_mask = (1UL << 12) - 1;

/* 用于储存早期内存分配器的一些信息。 */
static struct early_allocator early_allocator = { 0 };

/* 在前期内存分配系统中将页号转换为内存地址的宏 */
#define __PAGE_TO_ADDR(page_num)    \
    ((unsigned long)((page_num) << __page_shift))

/* 在前期内存分配系统中将内存地址转换为页号的宏 */
#define __ADDR_TO_PAGE(addr)    \
    ((unsigned long)(addr) >> __page_shift)

/* 将一个范围内的页加入伙伴系统。仅设置标志位。 */
static inline void __add_pages_to_buddy_system(
    unsigned long physical_start,
    unsigned long pfn_length
)
{
    struct page_struct *page = physical_to_page(physical_start);
    if (page == NULL) {
        return;
    }
    for (int i = 0; i < pfn_length; i++) {
        SetPageBuddy(page + i);
    }
}

/* 将一个范围内的页移出伙伴系统。仅设置标志位。 */
static inline void __del_pages_from_buddy_system(
    unsigned long physical_start,
    unsigned long pfn_length
)
{
    struct page_struct *page = physical_to_page(physical_start);
    if (page == NULL) {
        return;
    }
    for (int i = 0; i < pfn_length; i++) {
        ClearPageBuddy(page + i);
    }
}

int early_alloc_init(
    unsigned long area_page_start,
    unsigned long area_page_end,
    unsigned long free_page_start,
    unsigned long free_page_end,
    unsigned long page_shift
)
{
    /* 检查输入参数的合法性。 */
    if (area_page_start >= area_page_end || 
        free_page_end <= free_page_start || 
        free_page_end > area_page_end || 
        free_page_start < area_page_start
    ) {
        return -1;
    }
    if (((area_page_end << page_shift) >> page_shift) != area_page_end) {
        return -1;
    }

    /* 对内存页大小信息进行设置。 */
    __page_shift = page_shift;
    __page_mask = (1UL << __page_shift) - 1;

    /* 初始化位图。 */
    /// 计算位图所占的字节数。
    size_t __bitmap_size = area_page_end - area_page_start;
    __bitmap_size = (__bitmap_size + 7) >> 3;
    
    /// 获得存放位图的内存区域。
    unsigned long pages_needed = (__bitmap_size + __page_mask) >> __page_shift;
    /// 判断是否有足够的空间存放位图。
    if (free_page_end - free_page_start < pages_needed) {
        return -1;
    }
    /// 得到存放色位图所在的内存区域首地址。
    void *buffer = PHY_TO_VIRT(__PAGE_TO_ADDR(free_page_start));

    /// 初始化位图结构体。
    struct bitmap *__bitmap = bitmap_init(
        &(early_allocator.bitmap), 
        buffer, 
        __bitmap_size
    );

    /// 检查是否成功初始化。
    if (bitmap_size(__bitmap) < __bitmap_size) {
        return -1;
    }

    /// 将不可用的部分置位。
    bitmap_set_all(__bitmap);
    bitmap_clear_block(
        __bitmap, 
        free_page_start + pages_needed - area_page_start,
        free_page_end - area_page_start - 1
    );    

    /* 储存输入参数。 */
    early_allocator.page_start = area_page_start;
    early_allocator.pages_num = area_page_end - area_page_start;
    early_allocator.next_alloc_addr = VIRT_TO_PHY(buffer + __bitmap_size);
    early_allocator.next_page_offset =  \
        free_page_start + pages_needed - area_page_start;
    early_allocator.next_block_end =    \
        __PAGE_TO_ADDR(free_page_start + pages_needed);

    /// 初始化成功。
    return 0;
}

int early_pages_alloc_specific(
    unsigned long page_start, 
    unsigned long pages_num
)
{
    unsigned long area_start = early_allocator.page_start;
    unsigned long area_end = area_start + early_allocator.pages_num;
    unsigned long page_end = page_start + pages_num;
    
    /// 越界检测
    if (area_start >= page_end || 
        page_start < area_start || 
        page_end > area_end ||
        pages_num > early_allocator.pages_num
    ) {
        return -1;
    }

    /// 确保对应的内存页里面没有已被分配的内存页。
    if (bitmap_check_any_set(
            &(early_allocator.bitmap), 
            page_start - area_start, 
            page_end - area_start - 1)
    ) {
        return -1;
    }

    /// 进行分配。
    bitmap_set_block(
        &(early_allocator.bitmap), 
        page_start - area_start, 
        page_end - area_start - 1
    );

    /// 此处为了保险起见，检查被置位的整个区域是否真的全被置位。
    /// 如果返回false，说明这是一个bug。
    if (!bitmap_check_all_set(
            &(early_allocator.bitmap), 
            page_start - area_start, 
            page_end - area_start - 1)
    ) {
        /// @todo 检测bug。
        while (1);
    }

    /// 分配成功。
    __del_pages_from_buddy_system(
        __PAGE_TO_ADDR(page_start), 
        pages_num
    );
    return 0;
}

void *early_pages_alloc(unsigned long order, int order_aligned)
{
    /// 判断输入参数合法性。
    unsigned long pages_needed = 1UL << order;
    unsigned long order_mask = pages_needed - 1;
    if (pages_needed > early_allocator.pages_num) {
        return NULL;
    }

    /// 遍历当前每一个连续的空闲区域。
    unsigned long pos = 0;
    unsigned long page_start = 0;
    unsigned long page_start_aligned = 0;
    unsigned long page_end = 0;
    unsigned long pages_num = 0;
    unsigned long pages_num_aligned = 0;
    unsigned long area_start = early_allocator.page_start;
    unsigned long area_end = area_start + early_allocator.pages_num;
    unsigned long offset_start = 0;
    
    /// 遍历每个连续的空闲区域，将第一个满足条件的用于分配。
    for_each_early_free_area(pos, page_start, pages_num) {
        page_end = page_start + pages_num;
        /// 确保 page_end 不会超过整个指定内存区域的末尾。
        if (page_end > area_end) {
            page_end = area_end;
        }
        /// 如果指定按页号的阶数对齐，则按阶数对齐。
        if (order_aligned) {
            page_start_aligned = page_start + order_mask;
            page_start_aligned &= ~order_mask;
            pages_num_aligned = page_end - page_start_aligned;
        } else {
            page_start_aligned = page_start;
            pages_num_aligned = pages_num;
        }
        /// 如果当前区域过小，继续从下一个区域开始遍历。
        if (page_start_aligned > page_end || pages_num_aligned < pages_needed) {
            continue;
        }
        /// 找到合适的区域，进行分配并返回。
        offset_start = page_start_aligned - area_start;
        bitmap_set_block(
            &(early_allocator.bitmap), 
            offset_start, 
            offset_start + pages_needed - 1
        );
        /// 此处为了保险起见，检查被置位的整个区域是否真的全被置位。
        /// 如果返回false，说明这是一个bug。
        if (!bitmap_check_all_set(
                &(early_allocator.bitmap), 
                offset_start, 
                offset_start + pages_needed - 1)
        ) {
            /// @todo 检测bug。
            while (1);
        }
        
        __del_pages_from_buddy_system(
            __PAGE_TO_ADDR(page_start_aligned), 
            pages_needed
        );
        return PHY_TO_VIRT(__PAGE_TO_ADDR(page_start_aligned));
    }

    /// 没有找到合适的区域，返回NULL。
    return NULL;
}


void early_pages_free(void *page_beg, unsigned long nums)
{
    if (!bitmap_is_initialized(&(early_allocator.bitmap))) {
        /// 如果位图没有初始化，直接返回。
        return;
    }

    unsigned long page_start = __ADDR_TO_PAGE(VIRT_TO_PHY(page_beg));
    unsigned long page_end = page_start + nums;
    unsigned long area_start = early_allocator.page_start;
    unsigned long area_end = area_start + early_allocator.pages_num;
    if (page_start >= area_end || page_end <= area_start) {
        /// 计算出错。
        return;
    }
    /// 如果需要释放的区域末尾越界，强行让两者相等。
    if (page_end > area_end) {
        page_end = area_end;
    }
    if (page_end <= page_start) {
        return;
    }

    /// 进行释放。
    bitmap_clear_block(
        &(early_allocator.bitmap), 
        page_start - area_start, 
        page_end - area_start - 1
    );
    /// 此处为了保险起见，检查被置位的整个区域是否真的全被清除。
    /// 如果返回true，说明这是一个bug。
    if (bitmap_check_any_set(
            &(early_allocator.bitmap), 
            page_start - area_start, 
            page_end - area_start - 1)
    ) {
        /// @todo 检测bug。
        while (1);
    }
    __add_pages_to_buddy_system(VIRT_TO_PHY(page_beg), nums);
}


void *early_alloc(size_t size)
{
    unsigned long pages_needed = (size + __page_mask) >> __page_shift;
    if (pages_needed > 32 || pages_needed == 0) {
        /// 规定 early_alloc 分配的内存块不能大于32个内存页，且 size 不能为0。
        return NULL;
    }

    /// 获得当前剩余用于分配小内存块的内存空间，并按八字节对齐。
    unsigned long next_alloc_addr = early_allocator.next_alloc_addr;
    unsigned long next_block_end = early_allocator.next_block_end;
    void *result = NULL;
    next_alloc_addr += 7;
    next_alloc_addr &= ~7UL;

    /// 先尝试通过分配预留区域往后紧接着的内存页。
    if (next_block_end < next_alloc_addr + size) {
        unsigned long __new_alloc_end = next_alloc_addr + size;
        __new_alloc_end += __page_mask;
        __new_alloc_end &= ~__page_mask;
        if (early_pages_alloc_specific(
                next_block_end >> __page_shift, 
                (__new_alloc_end >> __page_shift) 
                - (next_block_end >> __page_shift)
            ) == 0
        ) {
            /// 分配成功，修改相关信息。
            next_block_end = __new_alloc_end;
            early_allocator.next_block_end = __new_alloc_end;
        }
    }

    /// 如果剩余空间足够，直接分配。
    if (next_block_end >= next_alloc_addr + size) {
        result = PHY_TO_VIRT(next_alloc_addr);
        next_alloc_addr += size;
        early_allocator.next_alloc_addr = next_alloc_addr;
        return result;
    }

    /// 剩余空间不够时，先判判断剩余空间是否大于单个内存页，如果大于则先把能释放的内存页都先
    /// 释放掉。
    early_alloc_block_compact();

    /// 重新申请内存页，若申请成功则进行分配。
    unsigned long order = 0;
    pages_needed -= 1;  /// 防止过多分配。
    while (pages_needed > 0) {
        order += 1;
        pages_needed >>= 1;
    }
    next_alloc_addr = VIRT_TO_PHY(early_pages_alloc(order, 0));
    if (next_alloc_addr == 0) {
        /// 申请内存页失败。
        return NULL;
    }

    /// 申请内存页成功，进行分配。
    /// 因为内存页的大小一定大于8字节，所以不用特意对齐。
    result = PHY_TO_VIRT(next_alloc_addr);
    next_block_end = next_alloc_addr;
    next_block_end += 1UL << (order + __page_shift);
    next_alloc_addr += size;
    early_allocator.next_alloc_addr = next_alloc_addr;
    early_allocator.next_block_end = next_block_end;
    return result;
}


void *early_pages_alloc_by_num(unsigned long num)
{
    /// 判断输入参数合法性。
    if (num == 0 || num > early_allocator.pages_num) {
        return NULL;
    }

    /// 遍历当前每一个连续的空闲区域。
    unsigned long pos = 0;
    unsigned long page_start = 0;
    unsigned long page_end = 0;
    unsigned long pages_num = 0;
    unsigned long area_start = early_allocator.page_start;
    unsigned long area_end = area_start + early_allocator.pages_num;
    unsigned long offset_start = 0;
    
    /// 遍历每个连续的空闲区域，将第一个满足条件的用于分配。
    for_each_early_free_area(pos, page_start, pages_num) {
        page_end = page_start + pages_num;
        /// 确保 page_end 不会超过整个指定内存区域的末尾。
        if (page_end > area_end) {
            page_end = area_end;
        }
        /// 如果当前区域过小，继续从下一个区域开始遍历。
        if (page_start > page_end || pages_num < num) {
            continue;
        }
        /// 找到合适的区域，进行分配并返回。
        offset_start = page_start - area_start;
        bitmap_set_block(
            &(early_allocator.bitmap), 
            offset_start, 
            offset_start + num - 1
        );
        /// 此处为了保险起见，检查被置位的整个区域是否真的全被置位。
        /// 如果返回false，说明这是一个bug。
        if (!bitmap_check_all_set(
                &(early_allocator.bitmap), 
                offset_start, 
                offset_start + num - 1)
        ) {
            /// @todo 检测bug。
            while (1);
        }
        
        __del_pages_from_buddy_system(
            __PAGE_TO_ADDR(page_start), 
            num
        );
        return PHY_TO_VIRT(__PAGE_TO_ADDR(page_start));
    }

    /// 没有找到合适的区域，返回NULL。
    return NULL;
}


int early_alloc_next_free_area(
    unsigned long last_pages_end,
    unsigned long *page_start,
    unsigned long *pages_num
)
{
    /// 检查输入参数的合法性。
    struct bitmap *bitmap = &(early_allocator.bitmap);
    unsigned long area_start = early_allocator.page_start;
    if (!bitmap_is_initialized(bitmap) ||
        page_start == NULL ||
        pages_num == NULL
    ) {
        return 0;
    }
    /// 保证 last_pages_end 不在整个指定内存区域之前。
    if (last_pages_end < early_allocator.page_start) {
        last_pages_end = early_allocator.page_start;
    }
    if (last_pages_end - area_start >= early_allocator.pages_num) {
        return 0;
    }

    /// 如果第 last_pages_end 位被置位，返回该位往后一个空闲内存区域的信息。
    unsigned long pos = last_pages_end - early_allocator.page_start;
    if (bitmap_check_bit(bitmap, pos)) {
        pos = bitmap_next_unset_block(bitmap, pos);
        /// 如果获得的下一个空闲区域的起始偏移为0，则说明已经遍历完成。
        if (pos == -1) {
            return 0;
        }
    }
    /// 获取当前进行遍历的内存区域的页数量。
    unsigned long num = bitmap_sizeof_unset_block(bitmap, pos);
    if (num > 0) {
        *page_start = early_allocator.page_start + pos;
        *pages_num = num;
        return 1;
    }
    return 0;
}


unsigned long early_alloc_block_compact(void)
{
    /// 先抛弃当前不完整的内存页。
    unsigned long next_alloc_addr = early_allocator.next_alloc_addr;
    unsigned long next_block_end = early_allocator.next_block_end;
    next_alloc_addr += __page_mask;
    next_alloc_addr &= ~__page_mask;
    unsigned long pages_to_free = next_block_end >> __page_shift;
    pages_to_free -= next_alloc_addr >> __page_shift;
    
    if (next_block_end <= next_alloc_addr || pages_to_free == 0) {
        /// 如果剩余空间不足一个完整的内存页，则不需要释放。
        return 0;
    }

    /// 进行内存页的释放。
    early_pages_free(PHY_TO_VIRT(next_alloc_addr), pages_to_free);
    early_allocator.next_block_end = next_alloc_addr;
    return pages_to_free;
}


int early_alloc_disable(void)
{
    early_allocator.pages_num = 0;
}


void __mark_earliest_allocated_pfn(void)
{
    unsigned long last_pfn_start = early_allocator.page_start;
    unsigned long pfn_end = early_allocator.pages_num + last_pfn_start;
    unsigned long pos;
    unsigned long page_start;
    unsigned long pages_num;
    for_each_early_free_area(pos, page_start, pages_num) {
        if (page_start - last_pfn_start > 0) {
            __del_pages_from_buddy_system(
                last_pfn_start << PAGE_SHIFT, 
                page_start - last_pfn_start);
        }
        last_pfn_start = page_start + pages_num;
    }
    if (pfn_end > last_pfn_start) {
        __del_pages_from_buddy_system(
            last_pfn_start << PAGE_SHIFT,
            pfn_end - last_pfn_start
        );
    }
}
